import uvicorn

from mcp.server.fastmcp import FastMCP
from mcp.server.sse import SseServerTransport

from starlette.requests import Request
from starlette.routing import Mount, Route
from starlette.applications import Starlette

mcp = FastMCP("auto code server...")


@mcp.tool()
def write_local_file(code: str, filename: str) -> str:
    """ 将代码写入到本地

    Args:
        code: 将要写入到本地文件的代码
        filename: 文件的名称
    """

    try:
        return f"成功写入：{filename}"
    except Exception as e:
        return f"失败：{str(e)}"

@mcp.tool()
def install_dependencies(code: str) -> str:
    """ 用于安装依赖包

    Args:
        code: 安装依赖包的命令
    """
    print('----> install code is:', code)
    return ''

@mcp.tool()
def checkout_dir(code: str) -> str:
    """ 用于切换工作目录

    Args:
        code: 切换工作目录的命令
    """
    print('checkout code is:', code)
    return ''

@mcp.tool()
def weather(address: str) -> str:
    """ 查询某个地区的天气情况

    Args:
        address: 地区名称
    """
    if address == '深圳':
        return '晴朗'
    elif address == '广州':
        return '暴雨'
    else:
        return '多云'


""" mcp server 启动服务器 """
class McpServer:
    @staticmethod
    def get_starlette_app(mcp_server, sse):
        async def handle_sse(request: Request) -> None:
            async with sse.connect_sse(
                    request.scope,
                    request.receive,
                    request._send,
            ) as (read_stream, write_stream):
                await mcp_server.run(
                    read_stream,
                    write_stream,
                    mcp_server.create_initialization_options(),
                )

        return Starlette(
            debug=True,
            routes=[
                Route("/sse", endpoint=handle_sse),
                Mount("/messages/", app=sse.handle_post_message),
            ]
        )

    def start(self, port: int):
        mcp_server, sse = mcp._mcp_server, SseServerTransport("/messages/")
        uvicorn.run(self.get_starlette_app(mcp_server, sse), host="0.0.0.0", port=port)


